package me.flyray.bsin.server.impl;

import static java.net.Proxy.Type.HTTP;
import static java.util.stream.Collectors.joining;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import me.flyray.bsin.exception.BusinessException;
import me.flyray.bsin.facade.service.LangchainChatService;
import me.flyray.bsin.server.biz.LangCahinBiz;
import me.flyray.bsin.utils.RespBodyHandler;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;

import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

//@Service
@Slf4j
public class LangchainChatServiceImpl implements LangchainChatService {

    @Autowired
    private LangCahinBiz langCahinBiz;
    @Value("${langchain4j.chat-model.openai.api-key}")
    private String OPENAI_API_KEY;
    @Value("${vector.datasource.host}")
    private String vectorHost;
    @Value("${vector.datasource.port}")
    private Integer vectorPort;

    @Autowired
    private EmbeddingModel embeddingModel;

    @Override
    public Map<String, Object> chat(Map<String, Object> requestMap) {
        String question = (String) requestMap.get("question");
        String answer;
        try {
            OpenAiChatModel model = OpenAiChatModel.builder()
                    .apiKey(OPENAI_API_KEY)
                    .proxy(new Proxy(HTTP, new InetSocketAddress("127.0.0.1", 8889)))
                    .build();
            answer = model.generate(question);
            System.out.println(answer);
        } catch (Exception e) {
            throw new BusinessException("100000", e.toString());
        }
        return RespBodyHandler.setRespBodyDto(answer);
    }

    @Override
    public Map<String, Object> chatWithDocument(Map<String, Object> requestMap) {

        String question = (String) requestMap.get("question");
        String collectionName = (String) requestMap.get("collectionName");
        // Specify the question you want to ask the model
        // String question = "Who is Charlie?";

        // 1.向量数据库
        EmbeddingStore<TextSegment> embeddingStore = MilvusEmbeddingStore.builder()
                .host(vectorHost)
                .port(vectorPort)
                .collectionName(collectionName)
                .dimension(512) // 384
                .build();

        // 2.Embed the question
        Embedding questionEmbedding = embeddingModel.embed(question).content();

        // 3.Find relevant embeddings in embedding store by semantic similarity
        // You can play with parameters below to find a sweet spot for your specific use case
        int maxResults = 3;
        double minScore = 0.7;
        List<EmbeddingMatch<TextSegment>> relevantEmbeddings
                = embeddingStore.findRelevant(questionEmbedding, maxResults, minScore);

        // 4.Create a prompt for the model that includes question and relevant embeddings

//        <指令>根据已知信息，简洁和专业的来回答问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题”，不允许在答案中添加编造成分，答案请使用中文。 </指令>
//<已知信息>{{ context }}</已知信息>、
//<问题>{{ question }}</问题>
        PromptTemplate promptTemplate = PromptTemplate.from(
                "Answer the following question to the best of your ability:\n"
                        + "\n"
                        + "Question:\n"
                        + "{{question}}\n"
                        + "\n"
                        + "Base your answer on the following information:\n"
                        + "{{information}}");

        String information = relevantEmbeddings.stream()
                .map(match -> match.embedded().text())
                .collect(joining("\n\n"));

        Map<String, Object> variables = new HashMap<>();
        variables.put("question", question);
        variables.put("information", information);

        // 5.generate prompt
        Prompt prompt = promptTemplate.apply(variables);

        ChatLanguageModel chatModel = OpenAiChatModel.builder()
                .apiKey(OPENAI_API_KEY)
                .proxy(new Proxy(HTTP, new InetSocketAddress("127.0.0.1", 8889)))
                .build();

        // 6.Send the prompt to the OpenAI chat model
        AiMessage aiMessage = chatModel.generate(prompt.toUserMessage()).content();

        // 7.See an answer from the model
        String answer = aiMessage.text();
        System.out.println(answer); // Charlie is a cheerful carrot living in VeggieVille...

        return RespBodyHandler.setRespBodyDto(requestMap);
    }


}

